//
// Created by SongpingWang on 2025/2/5.
//

#ifndef LSTM_CAPTCHA_CAPTCHA_RECOGNIZE_H
#define LSTM_CAPTCHA_CAPTCHA_RECOGNIZE_H
#include <torch/torch.h>
#include <torch/script.h>
#include "model_new.h"
#include "code_dataset_new.h"


class captcha_recognize
{
private:
    Model model{ nullptr };
    torch::nn::CTCLoss criterion;
    CUSTOM_INFO cust_info;

public:
    explicit captcha_recognize() : model(nullptr) {
        model = Model(cust_info.alphaNum.size());
    }
    ~captcha_recognize() = default;

    void Start(int BATCH_SIZE, int TOTAL_EPOCH,
               const std::string& train_dir, const std::string& valid_dir, const std::string& save_path);

    static std::pair<int, int> calculate_acc(const CUSTOM_INFO& cust_info,
                                             const torch::Tensor& output,
                                             const torch::Tensor& target,
                                             const torch::Tensor& target_lengths);
    void save_model(const std::string& save_path);
};

#endif //LSTM_CAPTCHA_CAPTCHA_RECOGNIZE_H
